#include "mnist.h"
#include <ctime>

#define nvalid 10000

vector< MATRIX > images1;
vector< MATRIX > images2;

vector< int > labels1;
vector< int > labels2;

MATRIX wr1, wr2, f1, f2, t1, t2, r1, r2, r3, f11, m11, w11;

void Random_Arr_real32( float *a, int numElements ) {

	srand( ( unsigned int )( time( 0 ) ) );

	for( int i = 0; i < numElements; ++i ) {

		a[ i ] = 1/( float )( rand(  ) + 1992 );
		if( rand(  ) & 1 ) { a[ i ] = -a[ i ]; }
	}
}

void readData(  ) {

	ReadMNISTimg( "train-images-idx3-ubyte", images1 );
	ReadMNISTlabel( "train-labels-idx1-ubyte", labels1 );

	ReadMNISTimg( "t10k-images-idx3-ubyte", images2 );
	ReadMNISTlabel( "t10k-labels-idx1-ubyte", labels2 );
}

void init( int nhide ) {

	float* twr1 = new float [nhide * 785];
	float* twr2 = new float [10 * (nhide+1) ];

	Random_Arr_real32( twr1, nhide * 785 );
	Random_Arr_real32( twr2, (nhide+1) * 10  );

	wr1.newMatrix( nhide, 785, twr1, RamMemory );
	wr2.newMatrix( 10, (nhide+1), twr2, RamMemory );

	f1.newMatrix( 1, nhide );
	f2.newMatrix( 1, 10  );

	t1.newMatrix( 1, 10  );
	t2.newMatrix( 1, nhide );

	r1.newMatrix( 1, nhide );
	r2.newMatrix( 10, (nhide+1) );
	r3.newMatrix( nhide, 785 );

	f11.newMatrix( 1, (nhide+1) );
	m11.newMatrix( 1, 785 );
	w11.newMatrix( 10, nhide );

	delete []twr1;
	delete []twr2;
}

void demo(  ) {

	float a[300] = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
	float b[300];
	float c[30000];

	for(int i = 0; i < 300; ++i) b[i] = a[i] = 1;

	MATRIX p( 1, 20, a );
	MATRIX q( 1, 10, b );
	MATRIX r( 20, 10, c );

	r.resulMulMatrix( p, q, MatrixT, Matrix );

	r.toArray( c );
	for( int i = 0; i < 200; ++i ) printf( "%.2f   ", c[i] );
}

float error(  ) {

	float resul = 0;

	int st = (int)images1.size() - nvalid;
	for( int i = 0; i < nvalid; ++i ) {

		float e [10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
		e[labels1[i+st]] = 1;

		//input layer
		f1.resulMulMatrix( images1[i+st], wr1, Matrix, MatrixT );
		f1.sigmoid();
		
		//output layer
		f2.resulMulMatrix( f1, wr2, Matrix, MatrixT ); //[1x10]
		f2.exp(  );
		f2.mulNumber( (float)1.0/f2.sum(  ) );
		f2.log(  );

		//error
		t1.copy( e );
		t1.mulMatrixDirect( f2 );

		resul += t1.sum();
	}

	return -resul;
}

void train( int nhide, float alpha, float lamda ) {

	int loop = (int)images1.size() - nvalid;
	for( int i = 0; i < loop; ++i ) {

		float e [10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
		e[labels1[i]] = 1;

		//input layer
		f1.resulMulMatrix( images1[i], wr1, Matrix, MatrixT );
		f1.sigmoid();
		
		//output layer
		f2.resulMulMatrix( f1, wr2, Matrix, MatrixT ); //[1x10]
		f2.exp(  );
		f2.mulNumber( (float)1.0/f2.sum(  ) );

		//error
		t1.copy( e );
		t1.subMatrix( f2 );
		
		//gradient
		w11.resulSubColum( wr2 );
		r1.resulMulMatrix( t1, w11, Matrix, Matrix );

		t2.resulSigmoidGradient( f1 );
		t2.mulMatrixDirect( r1 );

		f11.resulAddNewColum( f1, 1 );
		r2.resulMulMatrix( t1, f11, MatrixT, Matrix );

		m11.resulAddNewColum( images1[i], 1 );
		r3.resulMulMatrix( t2, m11, MatrixT, Matrix );

		r2.mulNumber( alpha );
		r3.mulNumber( alpha );

		//weight decay:
		wr2.mulNumber( lamda*alpha/(float)images1.size() + 1 );
		wr1.mulNumber( lamda*alpha/(float)images1.size() + 1 );

		//update weight
		wr2.addMatrix( r2 );		
		wr1.addMatrix( r3 );
	}
}

float test(  ) {

	int rate = 0;
	for( int i = 0; i < images2.size(); ++i ) {

		float e [10] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
		e[labels2[i]] = 1;

		f1.resulMulMatrix( images2[i], wr1, Matrix, MatrixT );
		f1.sigmoid();
		
		f2.resulMulMatrix( f1, wr2, Matrix, MatrixT ); //[1x10]
		f2.exp(  );
		f2.mulNumber( (float)1.0/f2.sum(  ) );

		if( f2.whereMax() == labels2[i] ) {

			++rate;
		}
	}

	return float(rate)/float(images2.size());	
}

int main(  ) {

	//read data
	readData(  );

	for( float lamda = 200; lamda < 3000; lamda *= 10 ) {

		for( float alpha = 0.003; alpha < 0.01; alpha *= 3  ) {

			for( int nhide = 100; nhide < 1000; nhide += 400 ) {

				float er1 = 1;
				float er2 = 0;

				//init
				init( nhide );

				while( abs( er1 - er2 ) > 0.5 ) {

					train( nhide, alpha, (float)1/lamda );
					er1 = er2;
					er2 = error(  );
				}

				printf( "nhide: %d; alpha: %.5f; lamda %.5f; accuracy: %.5f\n\n", nhide, alpha, 1/lamda, test(  ) );
			}
		}
	}

	return 0;
}
